// Package compiler generates bytecode for a Nix program.
package compiler

import (
	"bytes"
	"errors"
	"fmt"
	"slices"
	"strconv"
	"sync"

	"tvl.fyi/tvix/eval/builtins"
	"tvl.fyi/tvix/eval/internal/bytecode"
	"tvl.fyi/tvix/eval/internal/index"
	"tvl.fyi/tvix/eval/internal/op"
	p "tvl.fyi/tvix/eval/parser"
	"tvl.fyi/tvix/eval/value"
)

// builtinsAttrs is the set of built-in functions and values available in Nix.
var builtinsAttrs = sync.OnceValue(func() value.NixAttrs {
	m := value.NixAttrs{
		"null":  value.Null{},
		"true":  value.Bool(true),
		"false": value.Bool(false),
	}

	for k, v := range builtins.PureBuiltins() {
		m[value.String(k)] = v
	}

	return m
})

// toplevel is the set of identifiers available in the global scope (i.e.
// available not just at `builtins.foo`, but also directly at `foo`). This is
// not configurable, it is based on what Nix 2.3 exposed.
var toplevel = []value.String{
	"__curPos",
	"abort",
	"baseNameOf",
	"builtins",
	"derivation",
	"derivationStrict",
	"dirOf",
	"false",
	"fetchGit",
	"fetchMercurial",
	"fetchTarball",
	"fromTOML",
	"import",
	"isNull",
	"map",
	"null",
	"placeholder",
	"removeAttrs",
	"scopedImport",
	"throw",
	"toString",
	"true",
}

// globals is all global variables available in Nix. In the future this will
// become more configurable (custom builtins etc.)
var globals = sync.OnceValue(func() map[string]value.NixValue {
	b := builtinsAttrs()

	// Untie the knot of builtins referring to itself. They are wrapped in a
	// thunk to make cycle detection possible and avoid an actually infinite
	// data structure.
	bt := bytecode.NewEvaluatedThunk(b)
	b["builtins"] = bt

	// move the exposed builtins to the top-level scope
	g := make(map[string]value.NixValue)
	for k, v := range b {
		if slices.Contains(toplevel, k) {
			g[string(k)] = v
		}
	}

	return g
})

type Error []error

func (e Error) Error() string {
	var buf bytes.Buffer

	fmt.Fprintln(&buf, "errors while compiling:")

	for _, err := range e {
		fmt.Fprintln(&buf, "   ", err)
	}

	return buf.String()
}

// TODO: remove this function once things are stable
func (c *compiler) assertNodeType(n *p.Node, t p.NodeType) {
	if n.Type != t {
		c.error(fmt.Errorf("tvix bug: expected node type %q, but got %q", t, n.Type))
	}
}

// pushConstant pushes a literal value onto the current chunk and emits a
// corresponding instruction for it. This helper exists to ensure slots are
// always passed and initialised for these values.
func (c *compiler) pushConstant(slot index.Local, val value.NixValue) {
	ctx := c.unit()
	ctx.chunk.PushConstant(val)
	ctx.locals[slot].initialised = true
}

func Compile(p *p.Parser) (*bytecode.Chunk, error) {
	c := compiler{
		parser: p,
	}

	rootSlot := c.unit().declarePhantom()
	c.thunk(rootSlot, p.Result)

	// The final operation of any top-level Nix program must always be
	// `op.Force`. A thunk should not be returned to the user in an
	// unevaluated state (though in practice, a value *containing* a thunk
	// might be returned).
	c.emitForce()

	c.unit().chunk.PushOp(op.Return)

	if len(c.errors) > 0 {
		return nil, Error(c.errors)
	}

	return &c.unit().chunk, nil
}

type compiler struct {
	errors  []error
	parser  *p.Parser
	current *compilationUnit
}

func (c *compiler) error(err error) {
	if err != nil {
		c.errors = append(c.errors, err)
	}
}

func (c *compiler) unit() *compilationUnit {
	if c.current == nil {
		c.current = new(compilationUnit)
	}

	return c.current
}

func (c *compiler) pushUnit() {
	c.current = &compilationUnit{
		parent: c.current,
	}
}

func (c *compiler) popUnit() *compilationUnit {
	if c.current == nil {
		panic("tvix bug: compiler::popUnit called, but no compilationUnit exists")
	}

	popped := c.current
	c.current = popped.parent
	return popped
}

func (c *compiler) emitForce() {
	c.unit().chunk.PushOp(op.Force)
}

// compileNode compiles a single AST-node into corresponding instructions. The
// expectation is that this is always a self-contained Nix expression, which
// yields a value.
//
// compileNode is passed the value's expected stack slot, and marks the stack
// slot as initialised after compiling the value. This operation is important:
// When adding logic it is imperative that successful compilation does not exit
// early from this function.
func (c *compiler) compileNode(slot index.Local, n *p.Node) {
	switch n.Type {
	case p.IntNode:
		val, err := strconv.Atoi(c.parser.TokenString(n.Tokens[0]))
		if err != nil {
			c.error(fmt.Errorf("failed to parse integer literal: %w", err))
			return
		}

		c.unit().chunk.PushConstant(value.Int(val))

	case p.FloatNode:
		val, err := strconv.ParseFloat(c.parser.TokenString(n.Tokens[0]), 64)
		if err != nil {
			c.error(fmt.Errorf("failed to parse float literal: %w", err))
			return
		}

		c.unit().chunk.PushConstant(value.Float(val))

	case p.StringNode, p.IStringNode:
		c.compileString(slot, n)

	case p.TextNode:
		c.pushConstant(slot, value.String(c.parser.TokenString(n.Tokens[0])))

	case p.OpNode + '!':
		c.compileUnaryOp(n, op.Invert)

	case p.OpNode + p.Negate:
		c.compileUnaryOp(n, op.Negate)

	case p.OpNode + '+':
		c.compileBinOp(n, op.Add)

	case p.OpNode + '-':
		c.compileBinOp(n, op.Sub)

	case p.OpNode + '*':
		c.compileBinOp(n, op.Mul)

	case p.OpNode + '/':
		c.compileBinOp(n, op.Div)

	case p.OpNode + p.Concat:
		c.compileBinOp(n, op.Concat)

	case p.OpNode + p.Update:
		c.compileBinOp(n, op.Update)

	case p.OpNode + p.Eq:
		c.compileBinOp(n, op.Equal)

	case p.OpNode + p.And:
		c.compileAnd(n)

	case p.OpNode + p.Or:
		c.compileOr(n)

	case p.ListNode:
		// Open a temporary scope to correctly account for stack items
		// that exist during the construction.
		c.unit().beginScope()

		for _, elem := range n.Nodes {
			c.thunk(c.unit().declarePhantom(), elem)
		}

		c.unit().chunk.PushOpWithUint(op.List, uint64(len(n.Nodes)))
		c.unit().endScope()

	case p.ParensNode, p.InterpNode:
		c.compileNode(slot, n.Nodes[0])

	case p.SetNode:
		c.compileAttrs(slot, n.Nodes)

	case p.RecSetNode:
		if len(n.Nodes) == 0 {
			c.unit().chunk.PushConstant(value.NixAttrs{})
			break
		}

		c.error(errors.New("non-empty recursive sets not yet supported"))

	case p.SelectNode:
		ctx := c.unit()
		ctx.beginScope()

		// Push the set onto the stack
		c.compileNode(ctx.declarePhantom(), n.Nodes[0])
		// TODO: optimise_select if attrs are static

		// Compile each key fragment and emit access instructions.
		for _, fragment := range n.Nodes[1].Nodes {
			c.emitForce() // force set value
			c.compileAttr(ctx.declarePhantom(), fragment)
			ctx.chunk.PushOp(op.Select)
		}

		ctx.endScope()

	case p.LetNode:
		c.compileLet(slot, n)

	case p.IDNode:
		c.compileID(slot, n)

	case p.FunctionNode:
		c.compileFunction(slot, n)

	case p.ApplyNode:
		// To call a function, we leave its arguments on the stack,
		// followed by the function expression itself, and then emit a
		// call instruction. This way, the stack is perfectly laid out
		// to enter the function call straight away.
		c.unit().beginScope()
		c.thunk(c.unit().declarePhantom(), n.Nodes[1])
		c.compileNode(c.unit().declarePhantom(), n.Nodes[0])
		c.emitForce()
		c.unit().chunk.PushOp(op.Call)
		c.unit().endScope()

	case p.IfNode:
		c.compileIfElse(slot, n)

	default:
		c.error(fmt.Errorf("compiler does not yet support '%s'", n.Type.String()))
	}

	c.unit().locals[slot].initialised = true
}

func (c *compiler) compileUnaryOp(n *p.Node, op op.Op) {
	if len(n.Nodes) != 1 {
		c.error(fmt.Errorf("unary operation called with %d operand(s)", len(n.Nodes)))
		return
	}

	c.unit().beginScope()
	c.compileNode(c.unit().declarePhantom(), n.Nodes[0])
	c.emitForce()
	c.unit().chunk.PushOp(op)
	c.unit().endScope()
}

func (c *compiler) compileBinOp(n *p.Node, op op.Op) {
	if len(n.Nodes) != 2 {
		c.error(fmt.Errorf("binary operation called with %d operand(s)", len(n.Nodes)))
		return
	}

	c.unit().beginScope()
	c.compileNode(c.unit().declarePhantom(), n.Nodes[0])
	c.emitForce()
	c.compileNode(c.unit().declarePhantom(), n.Nodes[1])
	c.emitForce()
	c.unit().chunk.PushOp(op)
	c.unit().endScope()
}

func (c *compiler) compileAnd(n *p.Node) {
	c.unit().beginScope()

	// Both sides of the expression get compiled into the same slot, but we're
	// still accounting for it separately from the parent, just in case.
	slot := c.unit().declarePhantom()

	// Leave left-hand side value on the stack.
	c.compileNode(slot, n.Nodes[0])
	c.emitForce()

	// If this value is false, jump over the right-hand side - the whole
	// expression is false.
	jumpToEnd := c.unit().chunk.PushJump(op.JumpIfFalse)

	// Otherwise, remove the previous value and leave the right-hand side on the
	// stack. Its result is now the value of the whole expression.
	c.unit().chunk.PushOp(op.Pop)
	c.compileNode(slot, n.Nodes[1])
	c.emitForce()
	c.unit().chunk.PushOp(op.AssertBool)
	c.unit().chunk.PatchJump(jumpToEnd)

	c.unit().endScope()
}

func (c *compiler) compileOr(n *p.Node) {
	c.unit().beginScope()
	slot := c.unit().declarePhantom()
	c.compileNode(slot, n.Nodes[0])
	c.emitForce()

	// Opposite of compileAnd above: Short-circuit if lhs is true.
	jumpToEnd := c.unit().chunk.PushJump(op.JumpIfTrue)

	c.unit().chunk.PushOp(op.Pop)
	c.compileNode(slot, n.Nodes[1])
	c.emitForce()
	c.unit().chunk.PushOp(op.AssertBool)
	c.unit().chunk.PatchJump(jumpToEnd)

	c.unit().endScope()
}

func (c *compiler) compileAttr(slot index.Local, n *p.Node) {
	if id := c.staticIdent(n); id != "" {
		c.pushConstant(slot, value.String(id))
	} else {
		// TODO(tazjin): empty string is a valid identifier value, not dynamic ..
		c.error(fmt.Errorf("dynamic attribute selection not yet supported"))
	}
}

func (c *compiler) compileFunction(outerSlot index.Local, n *p.Node) {
	c.pushUnit()
	slot := c.unit().declarePhantom()
	c.unit().beginScope()

	// enter content (compiling actual function); what is throw idx?

	switch n.Nodes[0].Type {
	case p.IDNode:
		argname := c.parser.TokenString(n.Nodes[0].Tokens[0])
		idx, _ := c.unit().declareLocal(argname)
		c.unit().locals[idx].initialised = true
	default:
		c.error(fmt.Errorf("function argument type (%s) unsupported", n.Nodes[0].Type))
		c.popUnit()
		return
	}

	c.compileNode(slot, n.Nodes[1])

	pops := c.unit().endScope()
	c.unit().chunk.PushOpWithUint(op.CloseScope, uint64(pops))
	c.unit().chunk.PushOp(op.Return)

	// Compilation finished, construct the actual function value, remove the
	// context again, and push it onto the parent's constants.
	compiled := c.popUnit()
	f := function{
		lambda: &lambda{
			chunk: compiled.chunk,
		},
	}

	c.pushConstant(outerSlot, f)
}

// compileIfElse compiles conditional expressions using jumping instructions in
// the VM.
//
//	                      ┌────────────────────┐
//	                      │ 0  [ conditional ] │
//	                      │ 1   JUMP_IF_FALSE →┼─┐
//	                      │ 2  [  main body  ] │ │ Jump to else body if
//	                     ┌┼─3─←     JUMP       │ │ condition is false.
//	Jump over else body  ││ 4  [  else body  ]←┼─┘
//	if condition is true.└┼─5─→     ...        │
//	                      └────────────────────┘
func (c *compiler) compileIfElse(slot index.Local, n *p.Node) {
	// Slot note: Both bodies are compiled with the assumption that their
	// value ends up in `slot` (as only one of them is actually executed).
	// The condition is compiled with a slot in a temporary scope to avoid
	// affecting accounting in the bodies.
	c.unit().beginScope()
	c.compileNode(c.unit().declarePhantom(), n.Nodes[0]) // condition
	c.emitForce()
	jumpToElse := c.unit().chunk.PushJump(op.JumpIfFalse)

	// then-body
	c.unit().chunk.PushOp(op.Pop)   // discard condition value
	c.unit().endScope()             // end condition scope
	c.compileNode(slot, n.Nodes[1]) // compile then-branch
	jumpOverElse := c.unit().chunk.PushJump(op.Jump)

	// else-body
	c.unit().chunk.PatchJump(jumpToElse)
	c.unit().chunk.PushOp(op.Pop)             // discard condition value
	c.unit().locals[slot].initialised = false // reset initialisation state
	c.compileNode(slot, n.Nodes[2])           // compile else-branch
	c.unit().chunk.PatchJump(jumpOverElse)
}

// thunk compiles the given node in a new thunk. It emits the thunk as a
// constant.
//
// TODO: mark thunks that need capturing/"finalising"
func (c *compiler) thunk(outerSlot index.Local, n *p.Node) {
	// TODO: determine name (for errors) from outerSlot
	c.pushUnit()
	slot := c.unit().declarePhantom()
	c.unit().beginScope()
	c.compileNode(slot, n)

	pops := c.unit().endScope()
	if pops > 0 {
		c.unit().chunk.PushOpWithUint(op.CloseScope, uint64(pops))
	}

	compiled := c.popUnit()

	c.pushConstant(outerSlot, bytecode.NewSuspendedThunk(&compiled.chunk))
	c.unit().locals[outerSlot].captures = compiled.captures

	// TODO: unthunk (constants, builtins access, ...)? Worth it?
}
